#include "../inc/sched.h"

#define for_active_class_range(class, _from, _to)                  \
    for (const struct sched_class **f = _from, **t = _to, **c = f; \
         ({class = *c; c != t;}); c = next_active_class(c))

#define for_each_active_class(class) \
    for_active_class_range(class, __sched_class_highest(), __sched_class_lowest())

/*
 * Iterate only active classes. SCX can take over all fair tasks or be
 * completely disabled. If the former, skip fair. If the latter, skip SCX.
 */
static inline const struct sched_class **next_active_class(const struct sched_class **class)
{
    class++;

    return class;
}

static void prev_balance(struct rq *rq, struct task_struct *prev,
                         struct rq_flags *rf)
{
    const struct sched_class *start_class = prev->sched_class;
    const struct sched_class *class;

    /*
     * We must do the balancing pass before put_prev_task(), such
     * that when we release the rq->lock the task is in the same
     * state as before we took rq->lock.
     *
     * We can terminate the balance pass as soon as we know there is
     * a runnable task of @class priority or higher.
     */
    for_active_class_range(class, __sched_class_at(start_class), __sched_class_lowest())
    {
        if (class->balance && class->balance(rq, prev, rf))
            break;
    }
}

static struct task_struct *__pick_task(struct rq *rq, struct task_struct *prev,
                                       struct rq_flags *rf)
{
    struct task_struct *p = NULL;
    const struct sched_class *class;

    prev_balance(rq, prev, rf);

    for_each_active_class(class)
    {
        if (class->pick_next_task)
        {
            p = class->pick_next_task(rq, prev);
            if (p)
                break;
        }
        else
        {
            p = class->pick_task(rq);
            if (p)
            {
                put_prev_set_next_task(rq, prev, p);
                break;
            }
        }
    }

    return p;
}

/*
 * Pick up the highest-prio task:
 */
static inline struct task_struct *__pick_next_task(struct rq *rq, struct task_struct *prev,
                                                   struct rq_flags *rf)
{
    struct task_struct *p;

    p = __pick_task(rq, prev, rf);

    return p;
}

static inline struct task_struct *pick_next_task(struct rq *rq, struct task_struct *prev, struct rq_flags *rf)
{
    return __pick_next_task(rq, prev, rf);
}
